Top K frequent elements [Quick Select, Heap, Bucket Sort]

Time: O(N); Space: O(N); medium

Given a non-empty array of integers, return the k most frequent elements.

Example 1:

Input: nums = [1, 1, 1, 2, 2, 3], k = 2

Output: [1, 2]

Notes:

  • You may assume k is always valid, 1 <= k <= number of unique elements.

  • Your algorithm’s time complexity must be better than O(NLogN), where N is the array’s size.

1. Bucket Sort Solution

[3]:

import collections

class Solution1(object):
    def topKFrequent(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: List[int]
        """
        counts = collections.Counter(nums)
        buckets = [[] for _ in range(len(nums)+1)]

        for i, count in counts.items():
            buckets[count].append(i)

        result = []
        for i in reversed(range(len(buckets))):
            for j in range(len(buckets[i])):
                result.append(buckets[i][j])
                if len(result) == k:
                    return result

        return result

[4]:

s = Solution1()
nums = [1,1,1,2,2,3]
k = 2
assert s.topKFrequent(nums, k) == [1, 2]

2. Quick Select Solution

[5]:

import collections
from random import randint

class Solution2(object):
    """
    Time: O(N) ~ O(N^2), O(N) on average
    Space: O(N)
    """
    def topKFrequent(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: List[int]
        """
        counts = collections.Counter(nums)
        p = []
        for key, val in counts.items():
            p.append((-val, key))
        self.kthElement(p, k);

        result = []
        for i in range(k):
            result.append(p[i][1])
        return result


    def kthElement(self, nums, k):
        def PartitionAroundPivot(left, right, pivot_idx, nums):
            pivot_value = nums[pivot_idx]
            new_pivot_idx = left
            nums[pivot_idx], nums[right] = nums[right], nums[pivot_idx]
            for i in range(left, right):
                if nums[i] < pivot_value:
                    nums[i], nums[new_pivot_idx] = nums[new_pivot_idx], nums[i]
                    new_pivot_idx += 1

            nums[right], nums[new_pivot_idx] = nums[new_pivot_idx], nums[right]
            return new_pivot_idx

        left, right = 0, len(nums) - 1
        while left <= right:
            pivot_idx = randint(left, right)
            new_pivot_idx = PartitionAroundPivot(left, right, pivot_idx, nums)
            if new_pivot_idx == k - 1:
                return
            elif new_pivot_idx > k - 1:
                right = new_pivot_idx - 1
            else:  # new_pivot_idx < k - 1.
                left = new_pivot_idx + 1

[6]:

s = Solution2()
nums = [1,1,1,2,2,3]
k = 2
assert s.topKFrequent(nums, k) == [1, 2]

[7]:

import collections

class Solution3(object):
    """
    Time: O(Nlogk)
    Space: O(N)
    """
    def topKFrequent(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: List[int]
        """
        return [key for key, _ in collections.Counter(nums).most_common(k)]

[8]:

s = Solution3()
nums = [1,1,1,2,2,3]
k = 2
assert s.topKFrequent(nums, k) == [1, 2]

[9]:

class Solution4(object):
    def topKFrequent(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: List[int]
        """
        counts = {}
        for num in nums:
            if num not in counts:
                counts[num] = 0
            counts[num] += 1

        counts = [k_v for k_v in counts.items()]
        counts.sort(key=lambda x: x[1], reverse=True)

        return [counts[i][0] for i in range(k)]

[10]:

s = Solution4()
nums = [1,1,1,2,2,3]
k = 2
assert s.topKFrequent(nums, k) == [1, 2]